package com.mask.encryption.service.impl;

import com.mask.encryption.service.MaskEncryptionService;

import javax.crypto.Cipher;
import javax.crypto.KeyGenerator;
import javax.crypto.SecretKey;
import javax.crypto.spec.GCMParameterSpec;
import javax.crypto.spec.SecretKeySpec;
import java.nio.ByteBuffer;
import java.nio.charset.StandardCharsets;
import java.security.SecureRandom;
import java.util.Base64;

/**
 * AES-GCM 加密实现
 */
public class AesGcmEncryptionService implements MaskEncryptionService {

    private static final String TRANSFORMATION = "AES/GCM/NoPadding";
    private static final int DEFAULT_IV_LENGTH = 12; // 96-bit 推荐

    private final SecretKey secretKey;
    private final int gcmTagLength;
    private final SecureRandom secureRandom = new SecureRandom();

    public AesGcmEncryptionService(byte[] keyBytes, int gcmTagLength) {
        if (keyBytes == null || !(keyBytes.length == 16 || keyBytes.length == 24 || keyBytes.length == 32)) {
            // 自动生成 256-bit
            this.secretKey = generateAesKey(32);
        } else {
            this.secretKey = new SecretKeySpec(keyBytes, "AES");
        }
        this.gcmTagLength = gcmTagLength;
    }

    private SecretKey generateAesKey(int lengthBytes) {
        try {
            KeyGenerator kg = KeyGenerator.getInstance("AES");
            kg.init(lengthBytes * 8);
            return kg.generateKey();
        } catch (Exception e) {
            throw new IllegalStateException("AES key generation failed", e);
        }
    }

    @Override
    public byte[] encrypt(byte[] plaintext) {
        try {
            byte[] iv = new byte[DEFAULT_IV_LENGTH];
            secureRandom.nextBytes(iv);

            Cipher cipher = Cipher.getInstance(TRANSFORMATION);
            GCMParameterSpec spec = new GCMParameterSpec(gcmTagLength, iv);
            cipher.init(Cipher.ENCRYPT_MODE, secretKey, spec);
            byte[] cipherText = cipher.doFinal(plaintext);

            // 输出: IV | CIPHERTEXT
            ByteBuffer buffer = ByteBuffer.allocate(iv.length + cipherText.length);
            buffer.put(iv);
            buffer.put(cipherText);
            return buffer.array();
        } catch (Exception e) {
            throw new IllegalStateException("AES-GCM encrypt failed", e);
        }
    }

    @Override
    public byte[] decrypt(byte[] ciphertext) {
        try {
            ByteBuffer buffer = ByteBuffer.wrap(ciphertext);
            byte[] iv = new byte[DEFAULT_IV_LENGTH];
            buffer.get(iv);
            byte[] actualCipher = new byte[buffer.remaining()];
            buffer.get(actualCipher);

            Cipher cipher = Cipher.getInstance(TRANSFORMATION);
            GCMParameterSpec spec = new GCMParameterSpec(gcmTagLength, iv);
            cipher.init(Cipher.DECRYPT_MODE, secretKey, spec);
            return cipher.doFinal(actualCipher);
        } catch (Exception e) {
            throw new IllegalStateException("AES-GCM decrypt failed", e);
        }
    }

    @Override
    public String encryptToBase64(String plaintext) {
        byte[] encrypted = encrypt(plaintext.getBytes(StandardCharsets.UTF_8));
        return Base64.getEncoder().encodeToString(encrypted);
    }

    @Override
    public String decryptFromBase64(String base64Ciphertext) {
        byte[] cipher = Base64.getDecoder().decode(base64Ciphertext);
        return new String(decrypt(cipher), StandardCharsets.UTF_8);
    }
}


